package com.xrui.hbase.codec.impl;

import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Maps;

import com.xrui.hbase.DecodedCell;
import com.xrui.hbase.codec.CellEncoder;
import com.xrui.hbase.exception.EncodingException;
import com.xrui.hbase.schema.spec.CellSpec;
import org.apache.avro.AvroRuntimeException;
import org.apache.avro.Schema;
import org.apache.avro.generic.GenericContainer;
import org.apache.avro.io.DatumWriter;
import org.apache.avro.io.Encoder;
import org.apache.avro.io.EncoderFactory;
import org.apache.avro.specific.SpecificDatumWriter;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.util.Map;

/**
 * Serializes Avro cells to bytes for persistence in HBase.
 * <p>
 * <p>
 * An Avro cell encoder is specific to one column in a KijiTable.
 * Depending on the column specification, Avro cells embed the writer schema or not.
 * When embedded, the Avro schema ID/hash is prepended to the encoded value.
 * <p>
 */
public final class AvroCellEncoder implements CellEncoder {
    /**
     * Mapping from class names of Avro primitives to their corresponding Avro schemas.
     */
    public static final Map<String, Schema> PRIMITIVE_SCHEMAS;
    private static final Schema NULL_SCHEMA = Schema.create(Schema.Type.NULL);

    static {
        final Schema booleanSchema = Schema.create(Schema.Type.BOOLEAN);
        final Schema intSchema = Schema.create(Schema.Type.INT);
        final Schema longSchema = Schema.create(Schema.Type.LONG);
        final Schema floatSchema = Schema.create(Schema.Type.FLOAT);
        final Schema doubleSchema = Schema.create(Schema.Type.DOUBLE);
        final Schema stringSchema = Schema.create(Schema.Type.STRING);

        // Initialize primitive schema mapping.
        PRIMITIVE_SCHEMAS = ImmutableMap
            .<String, Schema>builder()
            .put(boolean.class.getCanonicalName(), booleanSchema)
            .put(Boolean.class.getCanonicalName(), booleanSchema)
            .put(int.class.getCanonicalName(), intSchema)
            .put(Integer.class.getCanonicalName(), intSchema)
            .put(long.class.getCanonicalName(), longSchema)
            .put(Long.class.getCanonicalName(), longSchema)
            .put(float.class.getCanonicalName(), floatSchema)
            .put(Float.class.getCanonicalName(), floatSchema)
            .put(double.class.getCanonicalName(), doubleSchema)
            .put(Double.class.getCanonicalName(), doubleSchema)
            .put(String.class.getCanonicalName(), stringSchema)
            .put(org.apache.avro.util.Utf8.class.getCanonicalName(), stringSchema)
            .put(java.nio.ByteBuffer.class.getCanonicalName(), Schema.create(Schema.Type.BYTES))
            .build();
    }

    /**
     * Specification of the column to encode.
     */
    private final CellSpec mCellSpec;

    /**
     * Cache of Avro DatumWriter.
     * <p>
     * <p>
     * Avro datum writers aren't thread-safe, but if we ensure the schema of a datum writer is not
     * modified, the datum writer becomes thread-safe.
     * </p>
     * <p>
     * <p>
     * This cache is not globally shared at present.
     * To share this map globally (ie. static) requires using a WeakIdentityHashMap:
     * a weak map is required to garbage collect unused schemas;
     * an identity map is also required as Schema.hashCode/equals are imperfect.
     * </p>
     */
    private final Map<Schema, DatumWriter<Object>> mCachedDatumWriters = Maps.newHashMap();

    /**
     * A byte stream for when encoding to a byte array.
     * <p>
     * Since we use the same instance for all encodings, this makes the encoder thread-unsafe.
     */
    private final ByteArrayOutputStream mByteArrayOutputStream = new ByteArrayOutputStream();

    /**
     * An encoder that writes to the above byte stream.
     */
    private final Encoder mByteArrayEncoder =
        EncoderFactory.get().directBinaryEncoder(mByteArrayOutputStream, null);

    /**
     * Configured reader schema for the column to encode.
     * <p>
     * This may currently be null if we only know the fully-qualified name of the record.
     * Eventually, this will always be populated so we can validate records being written against
     * the present reader schema.
     */
    private final Schema mReaderSchema;

    /**
     * Creates a new <code>KijiCellEncoder</code> instance.
     *
     * @param cellSpec Specification of the cell to encode.
     * @throws IOException on I/O error.
     */
    public AvroCellEncoder(final CellSpec cellSpec) throws IOException {
        mCellSpec = Preconditions.checkNotNull(cellSpec);
        Preconditions.checkArgument(cellSpec.isAvro());
        mReaderSchema = mCellSpec.getAvroSchema();
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public byte[] encode(final DecodedCell<?> cell) throws IOException {
        return encode(cell.getData());
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public synchronized <T> byte[] encode(final T cellValue) throws IOException {
        mByteArrayOutputStream.reset();

        // Get the writer schema for this cell.
        final Schema writerSchema = getWriterSchema(cellValue);

        // Encode the cell value:
        try {
            getDatumWriter(writerSchema).write(cellValue, mByteArrayEncoder);
        } catch (ClassCastException | AvroRuntimeException cce) {
            throw new EncodingException(cce);
        }
        return mByteArrayOutputStream.toByteArray();
    }

    /**
     * Gets a datum writer for a schema and caches it.
     * <p>
     * <p> Not thread-safe, calls to this method must be externally synchronized. </p>
     *
     * @param schema The writer schema.
     * @return A datum writer for the given schema.
     */
    private DatumWriter<Object> getDatumWriter(final Schema schema) {
        final DatumWriter<Object> existing = mCachedDatumWriters.get(schema);
        if (null != existing) {
            return existing;
        }
        final DatumWriter<Object> newWriter = new SpecificDatumWriter<>(schema);
        mCachedDatumWriters.put(schema, newWriter);
        return newWriter;
    }

    /**
     * Gets the writer schema of a specified value.
     *
     * @param <T>       is the java type of the specified value.
     * @param cellValue to create the Avro schema of.
     * @return an Avro schema representing the type of data specified.
     * @throws EncodingException if no Avro writer schema can be inferred from cellValue.
     */
    private <T> Schema getWriterSchema(final T cellValue) {
        if (cellValue == null) {
            return NULL_SCHEMA;
        } else if (cellValue instanceof GenericContainer) {
            return ((GenericContainer) cellValue).getSchema();
        } else if (mReaderSchema != null) {
            // Compute the writer schema using old semantics. This will only validate primitive schemas.
            return mReaderSchema;
        } else {
            final String className = cellValue.getClass().getCanonicalName();
            final Schema schema = PRIMITIVE_SCHEMAS.get(className);
            if (schema == null) {
                throw new EncodingException(String.format(
                    "Unable to infer Avro writer schema for value: '%s'", cellValue));
            }
            return schema;
        }
    }

    @Override
    public String getName() {
        return this.getClass().getName();
    }
}
